# 训练入口脚本
import yaml
from pathlib import Path
from src.data.loader import TextDataset, create_dataloader
from src.data.preprocess import load_raw_data, preprocess_texts
from src.model.bert_wrapper import get_bert_model
from src.train.trainer import train_model
from src.train.metrics import get_compute_metrics_fn
from src.transformers import TrainingArguments
from src.utils.logger import get_logger

from dotenv import load_dotenv
load_dotenv()  # 加载环境变量


def main():
    # 初始化日志
    logger = get_logger('train_script')
    logger.info('开始训练流程')

    # 加载配置
    with open('config/base.yaml', 'r') as f:
        base_config = yaml.safe_load(f)
    with open('config/train.yaml', 'r') as f:
        train_config = yaml.safe_load(f)
    with open('config/data.yaml', 'r') as f:
        data_config = yaml.safe_load(f)

    # 加载和预处理数据
    logger.info('加载并预处理数据')
    train_data = load_raw_data(data_config['train_path'])
    train_encodings = preprocess_texts(
        train_data['text'],
        base_config['model_name'],
        data_config['max_length']
    )
    train_dataset = TextDataset(train_encodings, train_data['label'].values)
    train_loader = create_dataloader(train_dataset, train_config['batch_size'])

    # 类似地加载验证集...

    # 加载模型
    logger.info(f'加载模型: {base_config['model_name']}')
    model = get_bert_model(
        base_config['model_name'],
        base_config['num_labels']
    )

    # 设置训练参数
    training_args = TrainingArguments(
        output_dir='models/finetuned/latest',
        learning_rate=train_config['learning_rate'],
        per_device_train_batch_size=train_config['batch_size'],
        num_train_epochs=train_config['num_epochs'],
        weight_decay=train_config['weight_decay'],
        fp16=train_config['fp16'],
        evaluation_strategy=train_config['evaluation_strategy'],
        save_strategy=train_config['save_strategy'],
    )

    # 获取评估指标函数
    compute_metrics = get_compute_metrics_fn(base_config['task_type'])

    # 开始训练
    logger.info('开始模型训练')
    train_model(model, train_dataset, None, training_args, compute_metrics)
    logger.info('训练完成')


if __name__ == '__main__':
    main()